import os
import re
import threading
import time
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import pandas as pd
import requests
from urllib.parse import urlparse
from ttkthemes import ThemedStyle
from concurrent.futures import ThreadPoolExecutor, as_completed
from tenacity import retry, stop_after_attempt, wait_fixed

headers = {
    "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
    "accept-language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
    "cache-control": "max-age=0",
    "if-modified-since": "Thu, 16 May 2024 23:50:05 GMT",
    "if-none-match": "\"c444356a2798944374fde31392b26384\"",
    "priority": "u=0, i",
    "sec-ch-ua": "\"Microsoft Edge\";v=\"137\", \"Chromium\";v=\"137\", \"Not/A)Brand\";v=\"24\"",
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": "\"Windows\"",
    "sec-fetch-dest": "document",
    "sec-fetch-mode": "navigate",
    "sec-fetch-site": "none",
    "sec-fetch-user": "?1",
    "upgrade-insecure-requests": "1",
    "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36 Edg/137.0.0.0"
}

class ImageDownloaderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图像下载工具 v2.0")
        self.root.geometry("800x600")
        self.root.minsize(600, 400)

        # 初始化变量
        self.file_path = ""
        self.save_dir = ""
        self.url_column = ""
        self.name_column = ""
        self.columns = []
        self.is_paused = False
        self.is_cancelled = False
        self.executor = None
        self.futures = []
        self.total_tasks = 0
        self.completed_tasks = 0

        # 初始化锁
        self.pause_lock = threading.Lock()
        self.cancel_lock = threading.Lock()

        self.create_widgets()
        self.style = ThemedStyle(self.root)
        self.style.set_theme("arc")  # 使用现代主题

    def create_widgets(self):
        # 文件选择
        self.file_label = ttk.Label(self.root, text="选择 Excel 文件:")
        self.file_label.pack(pady=10)
        self.file_button = ttk.Button(self.root, text="浏览", command=self.select_file)
        self.file_button.pack(pady=5)

        # 列选择
        self.column_frame = ttk.Frame(self.root)
        self.column_frame.pack(pady=10)

        self.url_label = ttk.Label(self.column_frame, text="URL 列:")
        self.url_label.grid(row=0, column=0, padx=5)
        self.url_combo = ttk.Combobox(self.column_frame, state="readonly")
        self.url_combo.grid(row=0, column=1, padx=5)

        self.name_label = ttk.Label(self.column_frame, text="图片名称列:")
        self.name_label.grid(row=1, column=0, padx=5)
        self.name_combo = ttk.Combobox(self.column_frame, state="readonly")
        self.name_combo.grid(row=1, column=1, padx=5)

        # 设置选项
        self.option_frame = ttk.Frame(self.root)
        self.option_frame.pack(pady=10)

        self.threads_label = ttk.Label(self.option_frame, text="线程数量:")
        self.threads_label.grid(row=0, column=0, padx=5)
        self.threads_entry = ttk.Entry(self.option_frame, width=10)
        self.threads_entry.insert(0, "5")
        self.threads_entry.grid(row=0, column=1, padx=5)

        self.timeout_label = ttk.Label(self.option_frame, text="超时时间 (秒):")
        self.timeout_label.grid(row=0, column=2, padx=5)
        self.timeout_entry = ttk.Entry(self.option_frame, width=10)
        self.timeout_entry.insert(0, "10")
        self.timeout_entry.grid(row=0, column=3, padx=5)

        # 保存路径选择
        self.save_button = ttk.Button(self.root, text="选择保存路径", command=self.select_save_dir)
        self.save_button.pack(pady=10)

        # 控制按钮
        self.control_frame = ttk.Frame(self.root)
        self.control_frame.pack(pady=10)

        self.start_button = ttk.Button(self.control_frame, text="开始下载", command=self.start_download)
        self.start_button.pack(side=tk.LEFT, padx=5)

        self.pause_button = ttk.Button(self.control_frame, text="暂停", command=self.toggle_pause)
        self.pause_button.pack(side=tk.LEFT, padx=5)

        self.resume_button = ttk.Button(self.control_frame, text="继续", command=self.resume_download)
        self.resume_button.pack(side=tk.LEFT, padx=5)
        self.resume_button.config(state=tk.DISABLED)

        self.cancel_button = ttk.Button(self.control_frame, text="取消", command=self.cancel_download)
        self.cancel_button.pack(side=tk.LEFT, padx=5)

        # 进度条和日志
        self.progress = ttk.Progressbar(self.root, orient="horizontal", length=400, mode="determinate")
        self.progress.pack(pady=10)

        self.log_text = tk.Text(self.root, height=10, wrap=tk.WORD)
        self.log_text.pack(pady=10, fill=tk.BOTH, expand=True)

    def log_message(self, message):
        self.log_text.insert(tk.END, f"{message}\n")
        self.log_text.see(tk.END)

    def select_file(self):
        self.file_path = filedialog.askopenfilename(filetypes=[("Excel files", "*.xlsx *.xls")])
        if self.file_path:
            self.load_columns()

    def load_columns(self):
        try:
            df = pd.read_excel(self.file_path)
            self.columns = list(df.columns)
            self.url_combo["values"] = self.columns
            self.name_combo["values"] = self.columns
        except Exception as e:
            self.log_message(f"读取 Excel 失败: {e}")

    def select_save_dir(self):
        self.save_dir = filedialog.askdirectory()
        if self.save_dir:
            self.log_message(f"图片将保存到: {self.save_dir}")

    def start_download(self):
        self.url_column = self.url_combo.get()
        self.name_column = self.name_combo.get()
        self.threads = self.threads_entry.get()
        self.timeout = self.timeout_entry.get()

        if not all([self.file_path, self.save_dir, self.url_column, self.name_column]):
            self.log_message("请确保所有选项已填写")
            return

        if not self.threads.isdigit() or int(self.threads) <= 0:
            self.log_message("线程数量必须为正整数")
            return

        if not self.timeout.isdigit() or int(self.timeout) <= 0:
            self.log_message("超时时间必须为正整数")
            return

        self.is_paused = False
        self.is_cancelled = False
        self.resume_button.config(state=tk.DISABLED)
        self.pause_button.config(state=tk.NORMAL)
        self.start_button.config(state=tk.DISABLED)
        self.cancel_button.config(state=tk.NORMAL)

        self.log_message("开始下载...")
        self.progress["value"] = 0
        self.completed_tasks = 0

        try:
            df = pd.read_excel(self.file_path)
            urls = df[self.url_column].dropna().tolist()
            names = df[self.name_column].dropna().tolist()
            self.total_tasks = len(urls)
            self.progress["maximum"] = self.total_tasks

            self.executor = ThreadPoolExecutor(max_workers=int(self.threads))
            self.futures = []

            for url, name in zip(urls, names):
                if self.is_cancelled:
                    break
                future = self.executor.submit(self.download_image, url, name, int(self.timeout))
                self.futures.append(future)

        except Exception as e:
            self.log_message(f"初始化下载失败: {e}")
            self.reset_controls()

    @retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
    def download_image(self, url, name, timeout):
        try:
            # 检查是否取消
            with self.cancel_lock:
                if self.is_cancelled:
                    return

            # 检查暂停
            while True:
                with self.pause_lock:
                    if not self.is_paused:
                        break
                if self.is_cancelled:
                    return
                time.sleep(0.5)

            # 检查 URL 是否为图片
            parsed_url = urlparse(url)
            if not parsed_url.scheme:
                self.log_message(f"[跳过] URL 缺少协议: {url}")
                return
            if not re.search(r'\.(jpg|jpeg|png|gif|bmp)$', parsed_url.path.lower()):
                self.log_message(f"[跳过] 不是图片格式: {url}")
                return

            # 下载图片
            with requests.Session() as session:
                session.headers.update(headers)
                with session.get(url, timeout=timeout, stream=True) as response:
                    response.raise_for_status()

                    # 检查是否取消
                    with self.cancel_lock:
                        if self.is_cancelled:
                            return

                    # 处理文件名
                    safe_name = re.sub(r'[<>:"/\\|?*\x00-\x1F]', '_', str(name))
                    ext = os.path.splitext(parsed_url.path)[1]
                    filename = f"{safe_name}{ext}"
                    filepath = os.path.join(self.save_dir, filename)

                    # 去重处理
                    counter = 1
                    while os.path.exists(filepath):
                        base, ext = os.path.splitext(filename)
                        filepath = os.path.join(self.save_dir, f"{base}_{counter}{ext}")
                        counter += 1

                    # 写入文件时检查取消状态
                    with open(filepath, "wb") as f:
                        for chunk in response.iter_content(chunk_size=8192):
                            if self.is_cancelled:
                                f.close()
                                os.remove(filepath)
                                return
                            f.write(chunk)

                    self.log_message(f"[成功] 下载完成: {filename}")
                    self.update_progress()

        except Exception as e:
            self.log_message(f"[失败] 下载失败: {url} - {e}")

    def update_progress(self):
        self.completed_tasks += 1
        self.progress["value"] = self.completed_tasks
        self.root.update_idletasks()

    def toggle_pause(self):
        with self.pause_lock:
            if not self.is_paused:
                self.is_paused = True
                self.pause_button.config(text="继续")
                self.resume_button.config(state=tk.NORMAL)
            else:
                self.is_paused = False
                self.pause_button.config(text="暂停")

    def resume_download(self):
        with self.pause_lock:
            self.is_paused = False
            self.pause_button.config(text="暂停")
            self.resume_button.config(state=tk.DISABLED)

    def cancel_download(self):
        with self.cancel_lock:
            self.is_cancelled = True
            self.is_paused = False
            self.log_message("下载已取消")
            for future in self.futures:
                future.cancel()
            self.reset_controls()

    def reset_controls(self):
        self.start_button.config(state=tk.NORMAL)
        self.pause_button.config(state=tk.DISABLED)
        self.resume_button.config(state=tk.DISABLED)
        self.cancel_button.config(state=tk.DISABLED)
        if self.executor:
            self.executor.shutdown(wait=False)
            self.executor = None
        self.futures.clear()

    def on_close(self):
        if self.executor and self.executor._work_queue.qsize() > 0:
            if messagebox.askokcancel("退出", "有任务正在运行，确定退出？"):
                self.executor.shutdown(wait=False)
                self.root.destroy()
        else:
            self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageDownloaderApp(root)
    root.protocol("WM_DELETE_WINDOW", app.on_close)
    root.mainloop()